const fs = require('fs');
const labelDef = require('./Def');
const BaseLabelLoader = require('../../BaseLabelLoader');

class COCOLabelLoader extends BaseLabelLoader {
    // 创建label对象
    create(bboxes, segs, labels, scores) {
        // 处理特别属性
        let result = [];
        for (let i = 0; i < bboxes.length; i++) {
            let jsonObj = labelDef();
            jsonObj.label_format = 'coco';
            jsonObj.id = i;
            let bbox = bboxes[i];
            let seg = segs ? segs[i] : [[]];
            let label = labels[i];
            let score = scores[i];

            if (bbox.length == 4) {
                let xmin = bbox[0];
                let ymin = bbox[1];
                let width = bbox[2];
                let height = bbox[3];
                jsonObj.bbox = [xmin, ymin, xmin + width, ymin + height];
            } else {
                jsonObj.bbox = [];
            }

            // 赋值通用属性
            let fields = [
                { name: 'category_id', value: 0 },
                { name: 'category', value: label },
                { name: 'area', value: 0 },
                { name: 'ignore', value: 0 },
                { name: 'iscrowd', value: 0 },
                { name: 'segmentation', value: seg },
                { name: 'score', value: score }
            ];
            for (let field of fields) {
                jsonObj[field.name] = field.value;
            }
            result.push(jsonObj);
        }
        return result;
    }

    // 加载数据
    async load(label_info) {
        // console.log("COCOLabelLoader::load label_info = ", label_info);
        try {
            const jsonData = fs.readFileSync(label_info.label_path, 'utf8');
            const cocoJson = JSON.parse(jsonData);
            // console.log(cocoJson);
            let result = [];
            for (let i = 0; i < cocoJson.length; i++) {
                let cocoAnno = cocoJson[i];

                // 处理特别属性
                let jsonObj = labelDef();
                jsonObj.label_format = 'coco';
                jsonObj.id = i; //并没有用原始标注数据中的id，考虑导出的时候再进行处理

                if (cocoAnno.bbox.length == 4) {
                    let xmin = cocoAnno.bbox[0];
                    let ymin = cocoAnno.bbox[1];
                    let width = cocoAnno.bbox[2];
                    let height = cocoAnno.bbox[3];
                    jsonObj.bbox = [xmin, ymin, xmin + width, ymin + height];
                } else {
                    jsonObj.bbox = [];
                }

                // 赋值通用属性
                let fields = [
                    { name: 'category_id', value: 0 },
                    { name: 'category', value: '' },
                    { name: 'area', value: 0 },
                    { name: 'ignore', value: 0 },
                    { name: 'iscrowd', value: 0 },
                    { name: 'segmentation', value: [[]] },
                    { name: 'num_keypoints', value: 0 },
                    { name: 'keypoints', value: [] }
                ];
                for (let field of fields) {
                    if (field.name in cocoAnno) {
                        jsonObj[field.name] = cocoAnno[field.name];
                    } else {
                        jsonObj[field.name] = field.value;
                    }
                }

                result.push(jsonObj);
            }
            // console.log(result);
            return result;
        } catch (error) {
            console.error(`Error load label ${label_info}, error: `, error);
            return { code: -1, msg: error.message };
        }
    }

    // 保存数据
    async save(label_info) {
        // console.log('COCOLabelLoader::save label_info = ', label_info);
        // 保存coco格式的json文件（忽略id，最后整体导出的时候再更新）
        let cocoJsonPath = label_info.label_url;
        let cocoJson = await this.fromStudio(label_info);
        // console.log(cocoJson);
        fs.writeFileSync(cocoJsonPath, JSON.stringify(cocoJson));

        // 更新Datasets和ListDB的信息（只更新下时间）
        let dataset_id = label_info.dataset_id;
        this.app.db && this.app.db.datasets.updateDataset({ id: dataset_id });
        let list_db = this.app.list_dbs[dataset_id];
        list_db && list_db.pairs.updateLabelPath(label_info.pairs_id);
    }

    async toString(label_info, categories = null) {
        let cocoJson = await this.fromStudio(label_info, categories);
        return JSON.stringify(cocoJson);
    }

    // 转换数据
    // 输入参数是labelOfDetetion类型
    async fromStudio(label_info, categories = null) {
        // console.log('COCOLabelLoader::fromStudio label_info = ', label_info);
        let dataset_id = label_info.dataset_id;
        let labels = label_info.labels;
        let cocoJson = [];
        if (!categories) {
            categories = await this.app.db.categories.getCategories(dataset_id);
        }
        for (let i = 0; i < labels.length; i++) {
            if (!labels[i]) continue;
            let jsonObj = labels[i].data;
            let cocoAnno = {
                bbox: []
            };

            // 处理category;
            let category = categories.find((el) => el.category_name === jsonObj.category);
            if (category) {
                jsonObj.category_id = category.category_id;
            }

            // 处理bbox
            if (jsonObj.bbox.length == 4) {
                let xmin = jsonObj.bbox[0];
                let ymin = jsonObj.bbox[1];
                let xmax = jsonObj.bbox[2];
                let ymax = jsonObj.bbox[3];
                cocoAnno.bbox[0] = xmin;
                cocoAnno.bbox[1] = ymin;
                cocoAnno.bbox[2] = xmax - xmin;
                cocoAnno.bbox[3] = ymax - ymin;
            } else {
                cocoAnno.bbox = [];
            }

            // 处理通用的属性
            if (jsonObj) {
                let fields = ['area', 'ignore', 'iscrowd', 'segmentation', 'num_keypoints', 'keypoints', 'category_id', 'category'];
                for (let field of fields) {
                    cocoAnno[field] = jsonObj[field];
                }
            }

            cocoJson.push(cocoAnno);
        }
        return cocoJson;
    }
}

module.exports = COCOLabelLoader;
